#include "color_correction/color_correction.hpp"
#include <iostream>
#include <string>

void printUsage() {
    std::cout << "颜色校正程序用法:" << std::endl;
    std::cout << "1. 校准模式: ./color_correction --calibrate --reference <参考图像> --input <输入图像> --output <JSON文件>" << std::endl;
    std::cout << "2. 应用模式: ./color_correction --apply --input <输入图像> --config <JSON文件> --output <输出图像>" << std::endl;
    std::cout << "3. 测试模式: ./color_correction --test --input <输入图像>" << std::endl;
    std::cout << std::endl;
    std::cout << "参数说明:" << std::endl;
    std::cout << "  --calibrate    : 校准模式，计算色彩校正参数" << std::endl;
    std::cout << "  --apply        : 应用模式，使用已保存的参数进行校正" << std::endl;
    std::cout << "  --test         : 测试模式，使用默认参数进行校正" << std::endl;
    std::cout << "  --reference    : 参考图像路径 (校准模式必需)" << std::endl;
    std::cout << "  --input        : 输入图像路径" << std::endl;
    std::cout << "  --config       : JSON配置文件路径" << std::endl;
    std::cout << "  --output       : 输出文件路径" << std::endl;
}

int main(int argc, char* argv[]) {
    std::cout << "=== OpenCV 颜色校正程序 ===" << std::endl;
    std::cout << "基于ArUco标记色卡的色彩校正矩阵实现" << std::endl;
    std::cout << std::endl;
    
    if (argc < 2) {
        printUsage();
        return -1;
    }
    
    std::string mode;
    std::string referencePath;
    std::string inputPath;
    std::string configPath;
    std::string outputPath;
    
    // 解析命令行参数
    for (int i = 1; i < argc; ++i) {
        std::string arg = argv[i];
        
        if (arg == "--calibrate") {
            mode = "calibrate";
        } else if (arg == "--apply") {
            mode = "apply";
        } else if (arg == "--test") {
            mode = "test";
        } else if (arg == "--reference" && i + 1 < argc) {
            referencePath = argv[++i];
        } else if (arg == "--input" && i + 1 < argc) {
            inputPath = argv[++i];
        } else if (arg == "--config" && i + 1 < argc) {
            configPath = argv[++i];
        } else if (arg == "--output" && i + 1 < argc) {
            outputPath = argv[++i];
        }
    }
    
    ColorCorrector corrector;
    
    if (mode == "calibrate") {
        // 校准模式
        if (referencePath.empty() || inputPath.empty()) {
            std::cout << "[错误] 校准模式需要指定参考图像和输入图像" << std::endl;
            printUsage();
            return -1;
        }
        
        if (outputPath.empty()) {
            outputPath = "color_correction_params.json";
        }
        
        std::cout << "[信息] 校准模式启动" << std::endl;
        std::cout << "[信息] 参考图像: " << referencePath << std::endl;
        std::cout << "[信息] 输入图像: " << inputPath << std::endl;
        std::cout << "[信息] 输出配置: " << outputPath << std::endl;
        
        // 加载图像
        cv::Mat refImage = cv::imread(referencePath);
        cv::Mat inputImage = cv::imread(inputPath);
        
        if (refImage.empty() || inputImage.empty()) {
            std::cout << "[错误] 无法加载图像文件" << std::endl;
            return -1;
        }
        
        std::cout << "[信息] 正在查找参考图像中的色卡..." << std::endl;
        cv::Mat refCard = corrector.findColorCard(refImage);
        
        std::cout << "[信息] 正在查找输入图像中的色卡..." << std::endl;
        cv::Mat inputCard = corrector.findColorCard(inputImage);
        
        if (refCard.empty() || inputCard.empty()) {
            std::cout << "[错误] 无法在图像中找到色卡" << std::endl;
            return -1;
        }
        
        // 显示检测到的色卡
        cv::imshow("Reference Color Card", refCard);
        cv::imshow("Input Color Card", inputCard);
        
        std::cout << "[信息] 正在计算色彩校正矩阵..." << std::endl;
        if (corrector.calculateColorCorrectionMatrix(refCard, inputCard)) {
            std::cout << "[信息] 色彩校正矩阵计算成功" << std::endl;
            
            // 保存参数
            if (corrector.saveParametersToJson(outputPath)) {
                std::cout << "[信息] 校准完成，参数已保存" << std::endl;
            }
            
            // 应用校正并显示结果
            cv::Mat correctedImage = corrector.applyColorCorrection(inputImage);
            corrector.showComparison(inputImage, correctedImage);
            
            // 保存校正后的图像
            std::string correctedOutputPath = "corrected_" + inputPath.substr(inputPath.find_last_of("/\\") + 1);
            cv::imwrite(correctedOutputPath, correctedImage);
            std::cout << "[信息] 校正后图像已保存: " << correctedOutputPath << std::endl;
            
        } else {
            std::cout << "[错误] 色彩校正矩阵计算失败" << std::endl;
            return -1;
        }
        
    } else if (mode == "apply") {
        // 应用模式
        if (inputPath.empty() || configPath.empty()) {
            std::cout << "[错误] 应用模式需要指定输入图像和配置文件" << std::endl;
            printUsage();
            return -1;
        }
        
        std::cout << "[信息] 应用模式启动" << std::endl;
        std::cout << "[信息] 输入图像: " << inputPath << std::endl;
        std::cout << "[信息] 配置文件: " << configPath << std::endl;
        
        // 加载配置
        if (!corrector.loadParametersFromJson(configPath)) {
            std::cout << "[错误] 无法加载配置文件" << std::endl;
            return -1;
        }
        
        // 加载图像
        cv::Mat inputImage = cv::imread(inputPath);
        if (inputImage.empty()) {
            std::cout << "[错误] 无法加载输入图像" << std::endl;
            return -1;
        }
        
        // 应用色彩校正
        std::cout << "[信息] 正在应用色彩校正..." << std::endl;
        cv::Mat correctedImage = corrector.applyColorCorrection(inputImage);
        
        // 显示结果
        corrector.showComparison(inputImage, correctedImage);
        
        // 保存结果
        if (outputPath.empty()) {
            outputPath = "corrected_" + inputPath.substr(inputPath.find_last_of("/\\") + 1);
        }
        cv::imwrite(outputPath, correctedImage);
        std::cout << "[信息] 校正后图像已保存: " << outputPath << std::endl;
        
    } else if (mode == "test") {
        // 测试模式 - 使用mine/env2.jpeg作为默认测试
        if (inputPath.empty()) {
            inputPath = "mine/env2.jpeg";
        }
        
        std::cout << "[信息] 测试模式启动" << std::endl;
        std::cout << "[信息] 测试图像: " << inputPath << std::endl;
        
        cv::Mat testImage = cv::imread(inputPath);
        if (testImage.empty()) {
            std::cout << "[错误] 无法加载测试图像: " << inputPath << std::endl;
            return -1;
        }
        
        // 显示原始图像
        cv::imshow("Test Image", testImage);
        
        // 尝试检测色卡
        std::cout << "[信息] 正在检测色卡..." << std::endl;
        cv::Mat colorCard = corrector.findColorCard(testImage);
        
        if (!colorCard.empty()) {
            cv::imshow("Detected Color Card", colorCard);
            std::cout << "[信息] 成功检测到色卡" << std::endl;
        } else {
            std::cout << "[警告] 未能检测到色卡" << std::endl;
        }
        
    } else {
        std::cout << "[错误] 未知模式: " << mode << std::endl;
        printUsage();
        return -1;
    }
    
    std::cout << std::endl;
    std::cout << "[信息] 按任意键关闭窗口..." << std::endl;
    cv::waitKey(0);
    cv::destroyAllWindows();
    
    return 0;
}
